package com.toxind.impl;

import java.io.File;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

@SuppressWarnings("unchecked")
public class ClassScanner {
//	private static final Log LOG = LogFactory.getLog(ClassScanner.class);

	private ClassLoader classLoader;

	public ClassScanner() {
		classLoader = ClassScanner.class.getClassLoader();
	}

	public ClassScanner(ClassLoader classLoader) {
		if (classLoader != null) {
			this.classLoader = classLoader;
		} else {
			classLoader = ClassScanner.class.getClassLoader();
		}
	}

	@SuppressWarnings("rawtypes")
	public Set<Class> scan(String clspass) {
		Set<Class> classes = new HashSet<Class>();
		loadClass(classes, classLoader, clspass);
		ClassLoader contextLoader = Thread.currentThread()
				.getContextClassLoader();
		if (!classLoader.equals(contextLoader))
			loadClass(classes, Thread.currentThread().getContextClassLoader(),
					clspass);
		return classes;
	}

	@SuppressWarnings("rawtypes")
	private void loadClass(Set<Class> classes, ClassLoader classLoader,
			String clspass) {
		List<URL> list = new ArrayList<URL>();
		ClassLoader cl = classLoader;
		while (cl != null) {
			if (cl instanceof URLClassLoader) {
				URL urls[] = ((URLClassLoader) cl).getURLs();
				list.addAll(Arrays.asList(urls));
			}
			cl = cl.getParent();
		}
		for (URL url : list) {
			System.out.println(url);
			File file;
			try {
				file = new File(url.toURI());
				if (file.getName().toLowerCase().endsWith(".jar")) {
					Set<String> fileNames = scannerJarFile(file);
					for (String fileName : fileNames) {
						if (fileName.toLowerCase().endsWith(".class")) {
							String className = fileName.replace(".class", "").replace(
									"/", ".");
							if (className.startsWith(clspass)) {
								try {
									Class clazz = classLoader.loadClass(className);
									classes.add(clazz);
								} catch (ClassNotFoundException e) {
//									LOG.error("class " + className + "can not be found");
								}
							}
						}
					}
				} else if (file.isDirectory()) {
					Set<String> fileNames = scannerDirectory(file,
							file.getAbsolutePath());
					for (String fileName : fileNames) {
						if (fileName.toLowerCase().endsWith(".class")) {
							String className = fileName.replace(".class", "").replace(
									File.separator, ".");
							if (className.startsWith(clspass)) {
								try {
									Class clazz = classLoader.loadClass(className);
									classes.add(clazz);
								} catch (ClassNotFoundException e) {
//									LOG.error("class " + className + "can not be found");
								}
							}
						}
					}
				} else {
					String fileName = file.getName();
					String className = fileName.replace(".class", "").replace("/",
							".");
					if (className.startsWith(clspass)) {
						try {
							Class clazz = classLoader.loadClass(className);
							classes.add(clazz);
						} catch (ClassNotFoundException e) {
//							LOG.error("class " + className + "can not be found");
						}
					}
				}
			} catch (URISyntaxException e1) {
				e1.printStackTrace();
			}

		}
	}

	public Set<String> scannerDirectory(File file, String root) {
		Set<String> names = new HashSet<String>();
		for (File f : file.listFiles()) {
			if (f.isDirectory()) {
				names.addAll(scannerDirectory(f, root));
			} else {
				String name = f.getAbsolutePath().substring(root.length() + 1);
				names.add(name);
			}
		}
		return names;
	}

	public Set<String> scannerJarFile(File file) {
		Set<String> names = new HashSet<String>();
		JarFile jarFile;
		try {
			jarFile = new JarFile(file);
			Enumeration<JarEntry> entries = jarFile.entries();
			while (entries.hasMoreElements()) {
				JarEntry entry = entries.nextElement();
				String name = entry.getName();
				names.add(name);
			}
		} catch (IOException e) {
//			LOG.error("file " + file.getName() + "can not be found");
		}
		return names;
	}

	@SuppressWarnings("rawtypes")
	public Set<Class> scan(String limitedPackage,
			Class<? extends Annotation> annotationCls) {
		Set<Class> classes = scan(limitedPackage);
		Set<Class> annotatedClasses = new HashSet<Class>();
		for (Class clazz : classes) {
			if (clazz.isAnnotationPresent(annotationCls))
				annotatedClasses.add(clazz);
		}
		return annotatedClasses;
	}
	
	public static void main(String[] args) {
		new ClassScanner().scan("junit");
	}
}
